# 脉冲累积下信噪比研究智能体
from typing import Dict, List
import json
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from apps.cpm.rlm_context import RlmContext
from apps.cpm.cpm_engine import CpmEngine
from core.radar_core import RadarCore
from common.app_manager import AppManager

class PaSnrAgent(object):
    def __init__(self):
        self.name = 'apps.cpm.agent.pa_snr_agent.PaSnrAgent'

    @staticmethod
    def run(context:RlmContext) -> Dict:
        
        '''
        作为整体任务第一步时，以用户自然语言为输入，输出为本步完成后的结果。
        '''
        observes = PaSnrAgent._observe(context=context)
        plans = PaSnrAgent._reason(observes=observes, context=context)
        return PaSnrAgent._act(observes=observes, context=context, plans=plans)
    
    @staticmethod
    def _observe(context:RlmContext) -> Dict:
        observes = {}
        prompt = PaSnrAgent._gen_params_parse_prompt(user_query=context.current_input)
        resp = CpmEngine.infer(query=prompt)
        # print(f'resp json: \n{resp[8:-3]};')
        params = json.loads(resp[8:-3])
        context.states['params'].update(params)
        observes.update(params)
        return observes
    
    @staticmethod
    def _reason(observes:Dict, context: RlmContext) -> List:
        plans = []
        return plans
    
    @staticmethod
    def _act(observes:Dict, context: RlmContext, plans: List) -> Dict:
        # 生成试验报告
         # 创建目录: work/reports/v001
        base_folder = 'work/reports/e005'
        Path(base_folder).mkdir(parents=True, exist_ok=True)
        # 创建图片目录：work/reports/v001/images
        images_folder = f'{base_folder}/images'
        Path(images_folder).mkdir(parents=True, exist_ok=True)
        # 执行业务逻辑
        c = 3e8 # 光速
        tau = float(observes['tau'])
        D_in = float(observes['D_in'])
        Theta_A = float(observes['Theta_A'])
        T_sc = float(observes['T_sc'])
        f_r = float(observes['f_r'])
        pt = float(observes['pt'])
        freq = float(observes['freq'])
        g = float(observes['G'])
        sigma = float(observes['sigma'])
        te = float(observes['te'])
        nf = float(observes['nf'])
        loss = float(observes['loss'])
        R_min = float(observes['R_min'])
        R_max = float(observes['R_max'])
        R_num = int(observes['R_num'])
        range_val = np.linspace(R_min, R_max, R_num)
        # 计算波长
        lambda_ = c / freq
        # 距离分辨率
        delta_R = (c*tau)/2.0
        # 带宽
        B = 1.0 / tau
        # 3dB波束宽度
        D = (D_in*2.54)/100.0
        theta_3dB = np.rad2deg(1.25 * (lambda_/D))
        # 目标照射时间
        dot_theta_scan = Theta_A / T_sc
        T_i = theta_3dB / dot_theta_scan
        n_P = int(np.ceil(theta_3dB / dot_theta_scan * f_r))
        # 计算单脉冲SNR
        snr1 = RadarCore.radar_eq(pt, freq, g, sigma, te, B, nf, loss, range_val)
        # 相参累积 (94个脉冲)
        n_pulses = 94
        Rnewci = (n_pulses**0.25) * range_val
        snrCI = snr1 + 10 * np.log10(n_pulses)
        # 绘图1：单脉冲 vs 相参累积
        fig, ax, axs = AppManager.init_matplotlib(rows=1, cols=1)
        range_km = range_val / 1000
        ax.plot(range_km, snr1, 'k', label='单脉冲')
        ax.plot(Rnewci / 1000, snr1, 'k-.', label='94脉冲相参累积')
        ax.axis([1, 12, -20, 45])
        ax.grid(True)
        ax.legend()
        ax.set_xlabel('探测距离 (km)')
        ax.set_ylabel('SNR (dB)')
        ax.set_title('相参累积效果')
        plt.savefig(f'{images_folder}/f00.png')
        # 非相参累积计算
        snr_b10 = 10**(snr1 / 10)  # 将SNR转换为线性值
        # 计算非相参累积后的SNR（根据雷达方程文献公式）
        SNR_1 = snr_b10 / (2 * n_pulses) + \
                np.sqrt((snr_b10**2) / (4 * n_pulses**2) + (snr_b10 / n_pulses))
        # 计算非相参累积损失
        LNCI = (1 + SNR_1) / SNR_1
        # 计算非相参累积增益
        NCIgain = 10 * np.log10(n_pulses) - 10 * np.log10(LNCI)
        # 计算等效距离
        Rnewnci = ((10**(0.1 * NCIgain))**0.25 * range_val)
        # 计算非相参累积后的SNR
        snrnci = snr1 + NCIgain
        # 绘图2：单脉冲 vs 非相参累积 vs 相参累积
        fig, ax, axs = AppManager.init_matplotlib(rows=1, cols=1)
        ax.plot(range_km, snr1, 'k', label='单脉冲')
        ax.plot(Rnewnci / 1000, snr1, 'k-.', label='94脉冲非相参累积')
        ax.plot(Rnewci / 1000, snr1, 'k:', label='94脉冲相参累积')
        plt.axis([1, 12, -20, 45])
        ax.grid(True)
        ax.legend()
        ax.set_xlabel('探测距离 (km)')
        ax.set_ylabel('SNR (dB)')
        ax.set_title('相参与非相参累积对比')
        plt.tight_layout()
        # plt.show()
        plt.savefig(f'{images_folder}/f01.png')
        # 更新上下文对象
        context.states['params'].update(observes)
        # plt.savefig(f'{images_folder}/f00.png')
        # 生成报告文本
        with open(f'{base_folder}/report.md', 'w', encoding='utf-8') as wfd:
            wfd.write(f'# 雷达设计：迷你设计案例1.2\n')
            wfd.write(f'# 1. 问题描述\n')
            wfd.write(f'{context.current_input}\n')
            wfd.write(f'# 2. 计算过程\n')
            wfd.write(f'## 2.1. 计算波长\n')
            wfd.write(f'$\\lambda = \\frac{{c}}{{f}}={lambda_}m$\n')
            wfd.write(f'## 2.2. 计算距离分辨率\n')
            wfd.write(f'$\\bigtriangleup R = \\frac{{c\\tau}}{{2}}={delta_R:0.2f}$\n')
            wfd.write(f'## 2.3. 带宽\n')
            wfd.write(f'$B=\\frac{{1}}{{\\tau}}={B}$\n')
            wfd.write(f'## 2.4. 3dB波束宽度\n')
            wfd.write(f'$\\theta_{{3dB}}={theta_3dB}\\circ$\n')
            wfd.write(f'## 2.5. 目标照射时间\n')
            wfd.write(f'$T_{{i}}={T_i}$\n')
            wfd.write(f'## 2.6. 得到脉冲数\n')
            wfd.write(f'$N_{{p}}={n_P}$\n')
            wfd.write(f'## 2.7. 信噪比与距离关系\n')
            wfd.write(f'单脉冲 vs 相参累积\n')
            wfd.write(f'![f00.png](./images/f00.png "")\n')
            wfd.write(f'单脉冲 vs 非相参累积 vs 相参累积\n')
            wfd.write(f'![f01.png](./images/f01.png "")\n')
        return {'result': 0, 'params': {}}
    
    @staticmethod
    def _gen_params_parse_prompt(user_query:str) -> str:
        template = '''
# 角色
假设你是资深雷达系统设计师

# 任务
你需要根据用户输入，提取出所有的参数，并转换为标准单位，以json格式输出。由于参数很多，而且参数表述和单位形式多样，你需要按照下面的步骤执行。

# 要求
你只需要给出json结果，不需要生成代码或进行解释。

# 步骤
1. 脉冲宽度（tau）：又叫脉宽，单位为s（秒）、ms（毫秒）、us（微秒）、ns（纳秒），若单位为ms时需乘以e-3，若单位为us时需乘以e-6，若单位为ns时需乘以e-9，提取值为v1；
2. 天线直径（D_in）：单位为in，提取值为v2；
3. 水平扫描范围（Theta_A）：单位为$\\circ$（度），若前面有$\\pm$号时需乘以2，提取值为v3；
4. 扫描时间（T_sc）：单位为s（秒）、ms（毫秒）、us（微秒）、ns（纳秒），若单位为ms时需乘以e-3，若单位为us时需乘以e-6，若单位为ns时需乘以e-9，提取值为v4；
5. 脉冲重复频率（f_r）：又叫PRF，单位为：Hz（赫兹）、KHz（千赫兹）、MHz（兆赫兹）、GHz（吉赫兹），若单位为KHz需乘以e3，若单位为MHz需乘以e6，若单位为GHz需乘以e9，提取值为v5；
6. 发射功率（pt）：又叫峰值功率、峰值发射功率，单位为：W（瓦）、KW（千瓦）、MW（兆瓦），若单位为KW时需乘以e3，若单位为MW时需乘以e6，提取值为v6；
7. 雷达频率（freq）：又叫中心频率、载频，单位为：Hz（赫兹）、KHz（千赫兹）、MHz（兆赫兹）、GHz（吉赫兹），若单位为KHz需乘以e3，若单位为MHz需乘以e6，若单位为GHz需乘以e9，提取值为v7；
8. 天线增益（G）：单位为dB（分贝），若只有数值需转为dB，公式为：公式为：dB=10*np.log10(G)，提取值为v8；
9. 目标截面积（sigma）：又叫目标RCS，单位为$m^{2}$（平方米），提取值为v9；
10. 噪声温度（te）：又叫工作温度、环境温度，单位为K（开尔文），若单位为摄氏度或度需加上275.12，提取值为v10；
11. 噪声系数（nf）：单位为dB（分贝），若只有数值需转为dB，公式为：公式为：dB=10*np.log10(nf)，提取值为v11；
12. 系统损耗（loss）：单位为dB（分贝），若只有数值需转为dB，公式为：公式为：dB=10*np.log10(nf)，提取值为v12；
13. 距离范围下限（R_min）：从"距离范围"附近的"从"得到，单位为：m（米）、Km（千米、公里），若单位为Km时需乘以e3，提取值为v13；
14. 距离范围上限（R_max）：从"距离范围"附近的"到"得到，单位为：m（米）、Km（千米、公里），若单位为Km时需乘以e3，提取值为v13；
15. 距离范围份数（R_num）：从"距离范围"附近的"分为"或"份"得到，提取值为v15；
11. 生成json结果：
{
    "tau": v1,
    "D_in": v2,
    "Theta_A": v3,
    "T_sc": v4,
    "f_r": v5,
    "pt": v6,
    "freq": v7,
    "G": v8,
    "sigma": v9,
    "te": v10,
    "nf": v11,
    "loss": v12,
    "R_min": v13,
    "R_max": v14,
    "R_num": v15
}

# 示例
输入：一个MMW雷达规范如下：中心频率$f=94GHz$，脉冲宽度$\\tau=50\\times 10^{-9} s$，峰值功率$P_{t}=4W$，水平扫描范围$\\Theta_{A}=\\pm 120^{\\circ}$，脉冲重复频率$PRF=10KHz$，噪声数$F=7dB$，天线直径为$D=12in$，天线增益$G=47dB$，目标RCS为$\\sigma=20m^{2}$，系统损失$L=10dB$，雷达扫描时间$T_{sc}=3s$，温度设为$T_{e}=290K$。距离范围为从1公里到12公里并分为1000份。请计算：波长$\\lambda$；距离分辨率$\\bigtriangleup R$；带宽$B$；天线半功率波束宽度；天线扫描率；计算SNR为10dB时的探测距离；绘制距离与SNR关系的曲线；计算目标得到的脉冲数，当使用脉冲累积时的探测距离，假设SNR不变（与单脉冲时相同）。
输出：
{
    "tau": 50e-9,
    "D_in": 12.0,
    "Theta_A": 240.0,
    "T_sc": 3.0,
    "f_r": 10e3,
    "pt": 4.0,
    "freq": 94e9,
    "G": 47.0,
    "sigma": 20.0,
    "te": 290,
    "nf": 7.0,
    "loss": 10.0,
    "R_min": 1e3,
    "R_max": 12e3,
    "R_num": 1000
}


# 当前输入
'''
        prompt = f'{template}{user_query}'
        return prompt